package com.ruoyi.service.impl;


import com.ruoyi.domain.entity.MyChatMessage;
import com.ruoyi.service.DouBaoAIService;
import com.volcengine.ark.runtime.model.completion.chat.ChatCompletionRequest;
import com.volcengine.ark.runtime.model.completion.chat.ChatMessage;
import com.volcengine.ark.runtime.model.completion.chat.ChatMessageRole;
import com.volcengine.ark.runtime.model.context.chat.ContextChatCompletionRequest;
import com.volcengine.ark.runtime.service.ArkService;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Service;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import javax.annotation.PostConstruct;
import javax.annotation.PreDestroy;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;

@Service
public class DouBaoAIServiceImpl implements DouBaoAIService {

    @Value("${douBao.apiKey}")
    private String apiKey; // 你的 API 密钥

    @Value("${douBao.AI_model_id}")
    private String AI_model_id; // 访问的模型ID

    // 初始化 ArkService
    private ArkService service;

    // 使用 @PostConstruct 确保 apiKey 注入后再初始化 ArkService
    @PostConstruct
    public void init() {
        if (apiKey == null || apiKey.isEmpty()) {
            throw new IllegalArgumentException("API Key cannot be null or empty");
        }
        // 使用注入的 apiKey 构建 ArkService 实例
        service = ArkService.builder().apiKey(apiKey).build();
    }

    // 在对象销毁时执行该方法
    @PreDestroy
    public void destroy() {
        // 调用 shutdownExecutor 来关闭 ArkService 资源
        if (service != null) {
            service.shutdownExecutor();
        }
    }


    public void getAIResponseStream(String question, SseEmitter emitter) {
        // 如果没有上下文ID，则创建一个新的上下文ID
//        String contextId = getContextId(service);  // 获取上下文ID的方法

        // 构建流式请求，使用上下文ID
        ChatCompletionRequest streamChatCompletionRequest = ContextChatCompletionRequest.builder()
//                .contextId(contextId) // 设置上下文ID
                .model(AI_model_id) // 替换为您自己的模型 ID
                .messages(Collections.singletonList(ChatMessage.builder().role(ChatMessageRole.USER)
                        .content(question).build()))
                .build();
        // 进行流式请求
        service.streamChatCompletion(streamChatCompletionRequest)
                .doOnError(Throwable::printStackTrace)
                .blockingForEach(
                        choice -> {
                            if (choice.getChoices().size() > 0) {
                                // 获取消息内容并通过 SseEmitter 逐步发送
                                try {
                                    emitter.send(choice.getChoices().get(0).getMessage().getContent());
                                } catch (Exception e) {
                                    emitter.completeWithError(e);
                                }
                            }
                        }
                );

        // 结束后关闭流
        emitter.complete();

    }

    @Override
    public void getAIMulAskStream(List<MyChatMessage> messageList, SseEmitter emitter) {
        ChatCompletionRequest streamChatCompletionRequest = buildStreamRequest(messageList);
        service.streamChatCompletion(streamChatCompletionRequest)
                .doOnError(Throwable::printStackTrace)
                .blockingForEach(
                        choice -> {
                            if (choice.getChoices().size() > 0) {
                                // 获取消息内容并通过 SseEmitter 逐步发送
                                try {
                                    emitter.send(choice.getChoices().get(0).getMessage().getContent());
                                } catch (Exception e) {
                                    emitter.completeWithError(e);
                                }
                            }
                        }
                );
        // 结束后关闭流
        emitter.complete();
    }

    // 根据对话列表，构建流式请求对象
    private ChatCompletionRequest buildStreamRequest(List<MyChatMessage> messageList){
        ArrayList<ChatMessage> chatList = new ArrayList<>();
        ChatMessage systemMsg = ChatMessage.builder().role(ChatMessageRole.SYSTEM).content("你是豆包，是一位爱好股票的行家").build();
        chatList.add(systemMsg);
        for (MyChatMessage message : messageList) {
            ChatMessageRole role = null;
            if(ChatMessageRole.USER.value().equals(message.getRole())){
                role = ChatMessageRole.USER;
            }else if(ChatMessageRole.ASSISTANT.value().equals(message.getRole())){
                role = ChatMessageRole.ASSISTANT;
            }
            ChatMessage chatMessage = ChatMessage.builder()
                    .role(role)
                    .content(message.getContent())
                    .build();
            chatList.add(chatMessage);
        }
        return ContextChatCompletionRequest.builder()
                .model(AI_model_id)
                .messages(chatList)
                .build();
    }

    public String getAIAsk(String question) {
        ArkService service = ArkService.builder().apiKey(apiKey).build();
        final List<ChatMessage> messages = new ArrayList<>();
        final ChatMessage systemMessage = ChatMessage.builder().role(ChatMessageRole.SYSTEM).content("你是豆包，是一位爱好股票的行家").build();
        final ChatMessage userMessage = ChatMessage.builder().role(ChatMessageRole.USER).content(question).build();
        messages.add(systemMessage);
        messages.add(userMessage);

        ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
                .model(AI_model_id)
                .messages(messages)
                .build();
        StringBuilder builder = new StringBuilder();
        service.createChatCompletion(chatCompletionRequest).getChoices().forEach(choice -> {
            builder.append(choice.getMessage().getContent());
        });
        service.shutdownExecutor();
        return builder.toString();
    }




    // 获取上下文ID的方法（无效，暂且不用）
//    private String getContextId(ArkService service) {
//        // 假设我们可以从缓存中获取上下文ID，使用本地变量或内存存储
//        String contextId = ContextCache.getChatContextId();
//        if (contextId == null) {
//            // 如果没有找到上下文ID，则创建新的上下文
//            CreateContextRequest createContextRequest = CreateContextRequest.builder()
//                    // 设置模型
//                    .model(AI_model_id)
//                    // 设置上下文模式为会话模式
//                    .mode(Const.CONTEXT_MODE_SESSION)
//                    // 设置系统消息
//                    .messages(Collections.singletonList(ChatMessage.builder().role(ChatMessageRole.SYSTEM)
//                            .content("你是股票行家AI人工助理").build()))
//                    // 设置上下文的生存时间（秒）
//                    .ttl(3600)
//                    .build();
//
//
//            // 创建上下文
//            CreateContextResult createContextResult = service.createContext(createContextRequest);
//
//            // 获取新创建的上下文ID
//            contextId = createContextResult.getId();
//
//            // 缓存上下文ID
//            ContextCache.setChatContextId(contextId);
//
//        }
//        System.out.println("created context, id = " + contextId);
//        return contextId;
//    }

}

